# SwingArmSuspensionEnv.py

import gymnasium as gym
from gymnasium import spaces
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

class SwingArmSuspensionEnv(gym.Env):
    """
    自定义摆臂式主动悬架Gym环境 (简化为四分之一车模型)
    """
    metadata = {'render_modes': ['human'], 'render_fps': 50}

    def __init__(self, render_mode=None):
        super().__init__()

        # --- 物理参数 ---
        self.m_s = 250.0  # 车身质量 (kg)
        self.m_u = 40.0   # 非簧载质量 (kg)
        self.k_s = 15000.0 # 悬架弹簧刚度 (N/m)
        self.k_t = 150000.0# 轮胎刚度 (N/m)
        self.c_s = 1000.0  # 悬架阻尼 (Ns/m)
        self.g = 9.81     # 重力加速度

        # --- 仿真参数 ---
        self.dt = 0.02  # 仿真步长 (s)
        self.max_steps = 500 # 每回合最大步数 (10秒仿真)

        # --- 状态空间和动作空间 ---
        # 状态: [x_s, dot_x_s, x_u, dot_x_u] (车身位移/速度, 车轮位移/速度)
        low_state = np.array([-0.3, -5.0, -0.15, -10.0], dtype=np.float32)
        high_state = np.array([0.3, 5.0, 0.15, 10.0], dtype=np.float32)
        self.observation_space = spaces.Box(low=low_state, high=high_state, dtype=np.float32)

        # 动作: 作动器力 Fa
        self.action_space = spaces.Box(low=-3000.0, high=3000.0, shape=(1,), dtype=np.float32)

        # --- 渲染相关 ---
        self.render_mode = render_mode
        self.fig = None
        self.ax = None
        
        # --- 内部状态 ---
        self.state = None
        self.time = 0
        self.current_step = 0

    def _system_dynamics(self, t, y, u, z_r_func):
        x_s, dot_x_s, x_u, dot_x_u = y
        z_r = z_r_func(t)
        
        F_k_s = self.k_s * (x_u - x_s)
        F_c_s = self.c_s * (dot_x_u - dot_x_s)
        F_t = self.k_t * (z_r - x_u)
        Fa = u[0]

        ddot_x_s = (F_k_s + F_c_s + Fa) / self.m_s
        ddot_x_u = (-F_k_s - F_c_s - Fa + F_t) / self.m_u
        return [dot_x_s, ddot_x_s, dot_x_u, ddot_x_u]

    def _get_road_profile(self, t):
        # 模拟一个减速带
        if 1.0 < t < 1.5:
            return 0.05 * (1 - np.cos(4 * np.pi * (t - 1.0)))
        return 0.0

    def step(self, action):
        sol = solve_ivp(
            fun=self._system_dynamics,
            t_span=[self.time, self.time + self.dt],
            y0=self.state,
            args=(action, self._get_road_profile),
            t_eval=[self.time + self.dt],
        )
        new_state = sol.y[:, -1]
        
        self.state = new_state
        self.time += self.dt
        self.current_step += 1

        x_s, _, x_u, _ = self.state
        ddot_x_s = self._system_dynamics(self.time, self.state, action, self._get_road_profile)[1]
        z_r_current = self._get_road_profile(self.time)

        # --- 奖励函数 ---
        w1, w2, w3, w4 = 1.0, 500.0, 500.0, 1e-7
        reward = -(
            w1 * ddot_x_s**2 +                           # 舒适性
            w2 * (x_s - x_u)**2 +                        # 悬架行程
            w3 * (x_u - z_r_current)**2 +                # 轮胎抓地力
            w4 * action[0]**2                            # 能量消耗
        )
        
        terminated = bool(abs(x_s) > 0.2 or abs(x_s - x_u) > 0.1)
        truncated = self.current_step >= self.max_steps
        
        if self.render_mode == 'human':
            self.render()

        return self.state, reward, terminated, truncated, {}

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self.state = np.zeros(self.observation_space.shape, dtype=np.float32)
        self.time = 0
        self.current_step = 0
        
        if self.render_mode == 'human' and self.fig is None:
            self._setup_render()

        return self.state, {}

    def _setup_render(self):
        self.fig, self.ax = plt.subplots(figsize=(10, 6))
        plt.ion() # 开启交互模式
        self.ax.set_xlim(-1, 1)
        self.ax.set_ylim(-0.3, 0.5)
        self.ax.set_aspect('equal')
        self.ax.set_title("Active Suspension Simulation")
        
        # 创建图形元素
        self.car_body = plt.Rectangle((-0.5, 0.2), 1.0, 0.1, fc='lightblue', ec='black')
        self.wheel = plt.Circle((0, 0), 0.05, fc='gray', ec='black')
        self.suspension = self.ax.plot([0, 0], [0.05, 0.2], 'k-', lw=4)[0]
        self.road = self.ax.plot([-1, 1], [0, 0], 'g-', lw=2)[0]
        self.ax.add_patch(self.car_body)
        self.ax.add_patch(self.wheel)
        self.info_text = self.ax.text(-0.95, 0.4, '', fontsize=10)

    def render(self):
        if self.render_mode != 'human':
            return

        x_s, _, x_u, _ = self.state
        z_r = self._get_road_profile(self.time)
        
        self.car_body.set_y(0.2 + x_s)
        self.wheel.set_center((0, z_r + x_u))
        self.suspension.set_data([0, 0], [self.wheel.center[1] + self.wheel.radius, self.car_body.get_y()])
        
        road_x = np.linspace(-1, 1, 200)
        road_y = [self._get_road_profile(self.time + (rx/2.0)) for rx in road_x] # 显示前方的路
        self.road.set_data(road_x, road_y)
        
        self.info_text.set_text(f"Time: {self.time:.2f}s\n"
                                f"Body Accel: {self._system_dynamics(self.time, self.state, [0], self._get_road_profile)[1]:.2f} m/s²")

        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
        plt.pause(0.001)

    def close(self):
        if self.fig is not None:
            plt.ioff()
            plt.close(self.fig)
            self.fig = None
